from mcp.server.fastmcp import FastMCP
from scipy.stats import norm, poisson, binom
import logging
import os
from typing import Dict, Literal, Union

# 配置日志
log_dir = os.path.join(os.path.dirname(__file__), 'logs')
os.makedirs(log_dir, exist_ok=True)
logging.basicConfig(
    filename=os.path.join(log_dir, 'distribution.log'),
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    encoding='utf-8'
)

logger = logging.getLogger(__name__)

# 创建MCP服务
mcp = FastMCP("Probability Distribution Server")

@mcp.tool()
def normal_distribution(
    mean: float,
    std_dev: float,
    x: float,
    calc_type: Literal['pdf', 'cdf'] = 'pdf'
) -> float:
    """正态分布计算"""
    if calc_type == 'pdf':
        return norm.pdf(x, loc=mean, scale=std_dev)
    else:
        return norm.cdf(x, loc=mean, scale=std_dev)

@mcp.tool()
def poisson_distribution(
    lambda_: float,
    x: int,
    calc_type: Literal['pdf', 'cdf'] = 'pdf'
) -> float:
    """泊松分布计算"""
    if calc_type == 'pdf':
        return poisson.pmf(x, mu=lambda_)
    else:
        return poisson.cdf(x, mu=lambda_)

@mcp.tool()
def binomial_distribution(
    n: int,
    p: float,
    x: int,
    calc_type: Literal['pdf', 'cdf'] = 'pdf'
) -> float:
    """二项分布计算"""
    if calc_type == 'pdf':
        return binom.pmf(x, n=n, p=p)
    else:
        return binom.cdf(x, n=n, p=p)

if __name__ == "__main__":
    logger.info("Probability distribution server starting")
    try:
        mcp.run(transport="stdio")
    except Exception as e:
        logger.exception("Server error")
        raise

if __name__ == "__main__":
    logger.info("Probability distribution server starting")
    try:
        mcp.run(transport="stdio")
    except Exception as e:
        logger.exception("Server error")
        raise